[DMD2 9/10] feat(dmd2): add the MiniMax H3 4-step DMD2 recipe with eval media and DMD2 metrics - #343

Ninth of the DMD2 stacked series, on top of #342. Adds the MiniMax H3 t2va 4-step DMD2 recipe for 2 nodes x 8 GPUs, and what is needed to judge a run by looking at it: eval videos at the right speed and with sound, kept on disk next to reference renders of the same prompts and seeds, plus distribution-matching and memory metrics.

What

• scripts/run_diffusion_dmd2_h3_t2va_16gpu_multinode.py: FastVideo's FastH3 Preview v1 Dense / Data-Free recipe (from checkpoint_metadata.json of FastVideo/FastVideo-FastH3-4-step-Preview-v1-Dense-DataFree) on --loss-type dmd2:
  • full-parameter student and critic from MiniMax-H3, frozen MiniMax-H3 teacher, fp32 master, bf16 forward; lr 2e-6 for both, AdamW (0.9, 0.999), weight decay 0.01, constant lr, grad clip 1.0;
  • 4 student steps (the uniform shift-12 grid, within 1.4e-4 in sigma of FastH3's rungs 999/749/500/250), stop_at_random_step backward simulation, ode rollout, unguided teacher;
  • --dmd2-critic-loss-space x0, --dmd2-score-sigma-range 0.001 0.999, --fsdp-flow-shift visual=5.0;
  • one rollout = 5 critic steps of 64 samples + 1 student step on the last 64, i.e. FastH3's generator update interval 5 at global batch 64; 200 rollouts = FastH3's 1000 steps;
  • prompts: the pinned vidprom_filtered_extended.txt (248,221 lines, 248,219 unique prompts), downloaded and split deterministically by the script into train / test / a 26-prompt eval_media set, with the output hashes pinned so every node builds the same files;
  • --reference teacher|fasth3 renders the eval set with MiniMax-H3 at 50 steps or with FastH3's student at 4 steps (the FastH3 checkpoint loaded by its Hub id on every node, with --diffusion-model-family h3 since its name matches no H3 pattern), through --debug-rollout-only eval.
• --rm-type none: a constant 0 reward with no scorer, for a loss that reads none. The rollout still asks for a reward, and DMD2 has no use for one. With it (and no --custom-rm-path), training requests also send rollout_return_generated_output=False, so the engine skips their VAE decode ([sglang-miles] [Diffusion] Let a rollout request skip the VAE decode sgl-project/sglang#46545); eval requests still decode. --rollout-video-dtype now applies to eval responses too.
• Media:
  • logged and saved videos play at --diffusion-fps (the hardcoded 8 stays the fallback); H3's config drops fps from the request, so the recipe sets 24 for playback only;
  • when the engine returns audio, it is muxed into the mp4 (AAC through imageio_ffmpeg's ffmpeg). Sample gains generated_audio / audio_sample_rate, read from the rollout response's generated_audio / audio_sample_rate when present;
  • --save-eval-media-dir: every eval sample's file under <dir>/eval_<rollout_id>/ with an index.jsonl of prompt, seed and reward, independent of --wandb-log-num-images; such samples keep their media past scoring;
  • RolloutManager.dispose joins the media threads, so the last eval (and a reference run, which ends right after it) is fully written and uploaded.
• Metrics:
  • DMD2 actor: dmd2_score_gap_abs_mean (mean |x0_fake − x0_real|), dmd2_teacher_dist_abs_mean (mean |x0 − x0_real|, g's normalizer), dmd2_x0_mean / dmd2_x0_std;
  • dmd2_grad_abs_mean, the two above and x0_mean_abs_diff are also logged per stop step as <key>_step{i}, through a STEP_SCHEMA family of new_metric_buffer(num_rollout_steps=...); a step no pair trained on logs nothing;
  • every trainer group: peak_memory_allocated_gb (this process, since the previous step) and device_memory_used_gb (the whole GPU, so sleeping groups and the engine count) per optimizer step.
• Docs: H3 model page (recipe, references, what to watch), CLI reference, rewards page, recipe verification (NV).

Why

• The goal of a first H3 run is to see whether the student learns. The fixed eval set renders at rollout 0 (the base model at 4 steps) and every --eval-interval with the same seeds, and the two reference runs render the same prompts and seeds, so files with the same name compare directly.
• The recipe follows the one H3 DMD2 recipe with a released checkpoint. The archived FastGen-DMD FLUX.3 recipe (rank-32 LoRA, lr 1e-4, batch 128) was superseded by its own authors (batch 64; full-parameter, Euler rollout and an x0 critic to match FastH3), and LoRA is out of scope for dmd2 here.
• Score shift: FastH3 sets score_timestep_shift 2.4 on its own clock, which FastGen-DMD documents as FastGen's score-time shift 5 (12 / 5 = 2.4); --fsdp-flow-shift visual=5.0 gives that density on the sigma axis. Worth re-checking against FastVideo's training code once it is published.
• Per-step metrics: DMD2 trains each sample at its random stop step, and a student that degrades usually does so at one step first; one mean over all steps hides that.
• Memory: the per-process peak and the device-wide use together show each group's headroom and what [DMD2 8/10] perf(ray): sleep each trainer group right after its train #342's earlier sleep saves.
• Upstream precedent for --rm-type none: miles' rm_hub has random / deterministic_random for runs that do not need a real reward.

Known gaps

• Audio needs [sglang-miles] [Diffusion] Return each sample's decoded audio in the rollout response sgl-project/sglang#46544, which adds generated_audio / audio_sample_rate to the rollout response; until it is in the image the saved videos are silent. The miles side here reads the fields when present.
• Skipping the training decode needs [sglang-miles] [Diffusion] Let a rollout request skip the VAE decode sgl-project/sglang#46545; an older engine drops the unknown field and decodes as before. On H3 (768p 4 s, 2×B200) it takes the decode from 1.27 s to 0 and the response from 396 MB to 64 MB per sample; at 5 s and 320 samples per rollout that is about 123 GB of decoded video per rollout not shipped.
• An image with #45230, #46544 and #46545 on top of sglang-miles is the branch zhihengy/sglang:dmd2/h3-dmd2-e2e.
• Only the video stream is distilled (the H3 rollout records no audio trajectory). FastH3 also mixes shapes; this recipe is 768p 16:9 only.

Tests

• tests/fast/backends/fsdp_utils/test_loss_hub_dmd2.py: the new metrics, overall and per stop step, on pairs with known teacher distance and score gap; steps no pair stopped at log nothing.
• tests/fast/rollout/test_eval_media.py: a video plays at the given fps (frame count and duration read back), the engine's audio ends up as a second stream with no leftover files, a single frame is a png, and _log_images with a save dir keeps every sample plus index.jsonl without W&B.
• tests/fast/rollout/test_rm_none.py: none gives 0 for every sample, globally and as a per-sample override.
• tests/fast/rollout/test_rollout_payload_media.py: training requests skip the decode only when no reward reads the media (none without a custom reward); eval requests always decode; the video dtype reaches eval requests.
• tests/fast/rollout/test_release_eval_sample_media.py: saved eval media are kept past scoring; released samples also drop their audio.
• tests/fast/utils/test_diffusion_rollout_response.py: the response's audio fields fill the sample, and are absent when the engine sends none.
• The recipe's train args and both --reference arg sets go through miles.utils.arguments.parse_args and its DMD2 validation offline (global batch 320, scoring range, critic loss space, x0 flow shift, rm none, eval media dir).
• The recipe has not run at its 16-GPU scale; [DMD2 e2e] test(e2e): 2-GPU H3 DMD2 proxy with a CI-recorded standard #366 runs a 2-GPU proxy of it as the e2e.

Checklist

• pre-commit hooks (black, isort, ruff at the pinned versions) pass on the changed files
• Added/updated tests for new behaviour (media fps / audio / save dir, rm none, DMD2 metrics, response audio)
• pytest tests/fast: 575 passed, 7 skipped; 4 tests fail in this environment and fail the same way on [DMD2 7/10] feat(dmd2): train the student with the distribution matching gradient from the critic group's scores #334 (test_hybrid_shard_mesh::test_mesh_axes and test_metric_buffer_dist::test_metric_buffer_reduction, whose torchrun children import another checkout, and two test_sft_h3_ref2va_encoder cases)
• If launch flags changed, python3 train_diffusion.py --help still parses (--save-eval-media-dir is listed)
• If a public flag was added, it appears in the CLI reference docs (--save-eval-media-dir; --rm-type none with its decode skip, eval --rollout-video-dtype, and the --diffusion-fps playback use in their rows)
• If an example was added, it has a real walkthrough (H3 model page, section 4.3)

Stack

DMD2: distill a flow-matching teacher into a few-step student with distribution matching. Merge in order:

1. [DMD2 1/10] [DMD2 1/10] refactor(ray): build each trainer's args from a trainer config list #335: per-trainer args from a trainer config list (merged)
2. [DMD2 2/10] [DMD2 2/10] feat(ray): train several colocated trainer groups in the sync driver #329: several colocated trainer groups in the sync driver (merged)
3. [DMD2 3/10] [DMD2 3/10] feat(ray): hand earlier trainer groups' outputs to later groups as external_data #330: hand earlier trainer groups' outputs to later groups as external_data
4. [DMD2 4/10] [DMD2 4/10] feat(fsdp): add a train-input hook that preprocesses train pairs before micro-batching #338: a train-input hook that preprocesses a trainer group's train pairs before micro-batching
5. [DMD2 5/10] [DMD2 5/10] feat(rollout): stop the rollout at a random step for DMD2 backward simulation #331: few-step rollout and the DMD2 step strategy (needs [Feat] diffusion-rl: stop a rollout at a requested step sgl-project/sglang#45230)
6. [DMD2 6/10] [DMD2 6/10] feat(dmd2): train a fake-score critic group from the teacher and score the student's samples #333: DMD2 critic group: critic training and teacher/critic scoring of its last step
7. [DMD2 7/10] [DMD2 7/10] feat(dmd2): train the student with the distribution matching gradient from the critic group's scores #334: --loss-type dmd2 on the actor
8. [DMD2 8/10] [DMD2 8/10] perf(ray): sleep each trainer group right after its train #342: sleep each trainer group right after its train
9. [DMD2 9/10] this PR: MiniMax H3 4-step DMD2 recipe, eval media and DMD2 metrics

On top of this PR as well: #366 (the 2-GPU e2e of this recipe), #367 and #368 (hook registry and loop state refactors), #372 (the teacher forward in the trunk).

A GAN head and other follow-ups come after this series.